#!/bin/bash

# 显式初始化conda（关键步骤）
# 注意：替换为你的conda安装路径下的conda.sh位置
source ~/ZYD/python/anaconda3/etc/profile.d/conda.sh  # 若用anaconda
# wandb agent 启动脚本
# 功能：在指定GPU上启动wandb agent，并可设置实验数量

# 环境名称
CONDA_ENV_NAME="sleepModel"

# 默认值
COUNT=""
GPU_ID=""

# 解析命令行参数
# 第一个参数为sweep ID
SWEEP_ID="$1"
shift  # 移除第一个参数（sweep ID）

# 解析其他选项
while [[ $# -gt 0 ]]; do
    case $1 in
        -c|--count)
            COUNT="$2"
            shift 2
            ;;
        -g|--gpu)
            GPU_ID="$2"
            shift 2
            ;;
        *)
            echo "未知参数: $1"
            echo "使用方法: $0 <sweep_id> [--count|-c <实验数量>] [--gpu|-g <GPU编号>]"
            exit 1
            ;;
    esac
done

# 检查必要参数
if [ -z "${SWEEP_ID}" ]; then
    echo "错误: 必须提供sweep ID"
    echo "使用方法: $0 <sweep_id> [--count|-c <实验数量>] [--gpu|-g <GPU编号>]"
    exit 1
fi

# 构建agent命令
AGENT_CMD="conda activate ${CONDA_ENV_NAME} && wandb agent ${SWEEP_ID}"

# 添加count参数（如果提供）
if [ -n "${COUNT}" ]; then
    AGENT_CMD="${AGENT_CMD} --count ${COUNT}"
fi

# 添加GPU限制（如果提供GPU ID）
if [ -n "${GPU_ID}" ]; then
    echo "使用GPU ${GPU_ID}"
    AGENT_CMD="CUDA_VISIBLE_DEVICES=${GPU_ID} ${AGENT_CMD}"
else
    echo "不限制GPU使用"
fi

# 打印GPU信息
if command -v nvidia-smi &> /dev/null; then
    echo -e "\n可用GPU信息:"
    nvidia-smi --query-gpu=index,name,memory.total --format=csv,noheader,nounits
fi

# 显示执行的命令
echo -e "\n执行命令: ${AGENT_CMD}"
echo "正在启动wandb agent..."

# 执行agent命令
eval "${AGENT_CMD}"